#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import asyncio
import argparse
import logging
import sys

# 添加项目根目录到路径
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from app.core.config import settings
from app.core.logging import setup_logging, get_logger
from app.repositories.milvus import milvus_client
from app.repositories.vector import VectorRepository

# 创建日志记录器
logger = get_logger("init_milvus")


async def initialize_milvus(project_id: str = None):
    """初始化Milvus
    
    Args:
        project_id: 项目ID（可选）
    """
    try:
        # 连接到Milvus
        await milvus_client.connect()
        
        # 创建向量仓库
        vector_repo = VectorRepository(milvus_client)
        
        if project_id:
            # 为指定项目创建集合
            logger.info(f"为项目 {project_id} 创建集合")
            await vector_repo.ensure_collections(project_id)
        else:
            # 为默认项目创建集合
            from app.core.utils import generate_uuid
            default_project_id = generate_uuid()
            logger.info(f"为默认项目创建集合，ID: {default_project_id}")
            await vector_repo.ensure_collections(default_project_id)
        
        logger.info("Milvus初始化完成")
    except Exception as e:
        logger.error(f"Milvus初始化失败: {str(e)}")
        raise
    finally:
        # 断开连接
        await milvus_client.disconnect()


async def main():
    """主函数"""
    # 设置日志
    setup_logging()
    
    # 解析命令行参数
    parser = argparse.ArgumentParser(description="初始化Milvus")
    parser.add_argument("--project-id", help="项目ID")
    args = parser.parse_args()
    
    # 初始化Milvus
    await initialize_milvus(args.project_id)


if __name__ == "__main__":
    # 运行异步主函数
    asyncio.run(main())